'''
Author: zx
Date: 2021-04-08 13:49:35
LastEditTime: 2022-04-12 09:34:13
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: /undefined/Users/zhangxin/Work/SlitlessSim/sls_lit_demo/simDemo.py
'''

import galsim
from SpecDisperser import *
# from numpy import *
import numpy as np
from scipy import interpolate
import astropy.constants as acon
import astropy.constants as C
import astropy.units as U
from astropy.table import Table
import math
from astropy.io import fits
import random
from pylab import *
import ipdb
import speclite
from speclite.filters import FilterResponse

U_SPEC_WAVE = U.erg/(U.s * U.cm**2 * U.AA)


# def generateDemoSED():
#     normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T,
#                           names=('WAVELENGTH', 'FLUX'))

'''
@description:
@param {*} fn: file name, include 2 column, wavelength(A)  and flux(erg/s/cm2/A)
@param {*} s: band start , unit A
@param {*} e; end, unit A
@param {*} deltL: sample interval for SED
@return {*} sed, unit photo/s/m2/A
'''
def generateSEDfromFiles(fn, s, e, deltL):
    """
    s: lambda start, unit A
    e: lambda end, unit A

    return:
    SEDs is array, 2-dim, (gal_num+1)*(wavelength size), last row is wavelength
    """
    lamb = np.arange(s, e + deltL, deltL)
    spec_orig = np.loadtxt(fn)
    speci = interpolate.interp1d(spec_orig[:, 0], spec_orig[:, 1])
    y = speci(lamb)
    # sed_unit  = U.erg/U.s/U.cm**2/U.AA
    # flux_unit = (1/U.s/U.m**2/U.AA).unit
    # cfactor = (1.0 * sed_unit / (C.h.unit * C.c.unit/ U.AA)).to(flux_unit).value
    # cfactor = 1e-13
    # erg/s/cm2/A --> photo/s/m2/A
    flux = y * lamb / (acon.h.value * acon.c.value) * 1e-13
    SED = Table(np.array([lamb, flux]).T,names=('WAVELENGTH', 'FLUX'))

    return SED

def addReadoutNois(img = None, readout = 5):
    import random
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            img[i,j] += round(random.gauss(mu = 0, sigma = readout))

    return img
'''
description: integrate the SED file in the range of the band
param {*} spec
param {*} throughput
param {*} aper: aperture of telescope
return {*}
'''
def intCSSTFlux(spec=None, throughput=None,aper = 2):
    # specFn='/Users/zhangxin/Work/ImageSim/Mutil-band/super_simu_test/sample/Imsim/spec_modi/gss10413.spec'
    # throughputFn = '/Users/zhangxin/Work/ImageSim/Mutil-band/super_simu_test/sample/Imsim/CSSOS_Throughput/i_throughput.txt'
    # aper = 2

    # spec = loadtxt(specFn)
    # thr = loadtxt(throughputFn)

    # spec in photo/s/m2/A

    ids = (spec['WAVELENGTH']>2000) & (spec['WAVELENGTH']<11000)

    spec_w = spec['WAVELENGTH'][ids]
    # spec_lam_phot = spec['FLUX'][ids1][ids2]*(3631*1e-30/cons.h.value/(spec_w)*1e4)
    spec_lam_phot = spec['FLUX'][ids]

    speci = interpolate.interp1d(spec_w, spec_lam_phot)
    thri = interpolate.interp1d(throughput['WAVELENGTH'], throughput['SENSITIVITY'])

    thr_slect_ids = throughput['SENSITIVITY'] > 0.001
    s_w = throughput['WAVELENGTH'][thr_slect_ids][0]
    e_w = throughput['WAVELENGTH'][thr_slect_ids][-1]

    if s_w<spec_w[0]:
        s_w = spec_w[0]+1
    if e_w>spec_w[-1]:
        e_w = spec_w[-1]-1
    w_i = np.linspace(np.ceil(s_w), np.floor(e_w), int((e_w - s_w)/0.1) + 1)

    flux_thr = speci(w_i)* thri(w_i)

    int_flux = np.trapz(flux_thr,w_i)

    # try to use the speclite package to get the same result
    band = FilterResponse(
        wavelength = throughput['WAVELENGTH'] * U.Angstrom,
        response   = throughput['SENSITIVITY'], meta=dict(group_name='CSST', band_name='temp')
      )
    wave = spec['WAVELENGTH'] * U.AA
    flux = spec['FLUX'] * U_SPEC_WAVE / (wave.value / (acon.h.value * acon.c.value) * 1e-13) # photo/s/m2/A -->  erg/s/cm2/A
    int_flux_new = band.convolve_with_array(
        wave, flux, photon_weighted=True,
        interpolate=True, axis=-1, units=flux.unit
      )
    int_flux_new = int_flux_new.si
    print(f'old result: {int_flux}')
    print(f'new result: {int_flux_new}') # unit: e-/m^2/s

    int_flux_aper = int_flux*math.pi*(aper/2)*(aper/2)
    return int_flux_aper # unit: photos

'''
@description:
@param {*} sedFile: SED file name
@param {*} specConfile: config file like aXe
@param {*} s_n: sersic index
@param {*} re: effective radii
@param {*} pa: position angle
@param {*} ell: ellipticity
@param {*} t: exposure time once
@param {*} aper: csst aperter, 2 meter
@param {*} xcenter: the center x of a source in CCD (direct image)
@param {*} ycenter: the center y of a source in CCD (direct image)
@param {*} img: ccd img
@param {*} psf: psf
@return {*}
'''
def generateSpec(sedFile = 'data/sed/sed_67084.txt', specConfile = 'data/conf/CSST_GI2.conf',s_n = 1.0, re = 0.1, pa = 90,ell = 0.6, t = 150, aper = 2.0, xcenter = 5000,ycenter = 5000, img = None, psf = None):
    sed = generateSEDfromFiles(sedFile,2000,10000,0.5) # in photo/s/m2/A

    gal = galsim.Sersic(s_n, half_light_radius=re)

    gal_pa = pa * galsim.degrees
    gal_ell = gal.shear(q=ell, beta=gal_pa)

    conv_gal = galsim.Convolve([gal_ell,psf])

    #       ==== normalized image profile (sum to 1) =======
    stamp = conv_gal.drawImage(wcs=galsim.PixelScale(p_size)) * t*math.pi*(aper/2)*(aper/2)

    # left-bottom corner position of the stamp image in the raw image
    origin_star = [ycenter - (stamp.center.y - stamp.ymin),
                            xcenter - (stamp.center.x - stamp.xmin)]

    sdp = SpecDisperser(orig_img=stamp, xcenter=xcenter,
                                        ycenter=ycenter, origin=origin_star,
                                        tar_spec=sed, # in photo/s/m2/A
                                        conf=specConfile,
                                        isAlongY=0)
    ipdb.set_trace()

    spec_orders = sdp.compute_spec_orders()
    for k, v in spec_orders.items():
        img_s = v[0]
        origin_order_x = v[1]
        origin_order_y = v[2]
        specImg = galsim.ImageF(img_s)

        specImg.setOrigin(origin_order_x, origin_order_y)

        bounds = specImg.bounds & img.bounds
        if bounds.area() == 0:
            continue
        img[bounds] = img[bounds] + specImg[bounds]

'''
description:
param {*} sedFile: SED file name
param {*} throughputFn: transmission of all systerm include (filter,optical systerm, CCD QE...)
param {*} s_n: sersic index
param {*} re: effective radii
param {*} pa: position angle
param {*} ell: ellipticity
param {*} t: exposure time once
param {*} aper: csst aperter, 2 meter
param {*} xcenter: the center x of a source in CCD (direct image)
param {*} ycenter: the center y of a source in CCD (direct image)
param {*} img: ccd img
param {*} psf: psf
return {*}
'''
def generateGal(sedFile = 'data/sed/sed_67084.txt', throughputFn = 'data/CSST_Throughput/i.Throughput.fits',s_n = 1.0, re = 0.1, pa = 90,ell = 0.6, t = 150, aper = 2.0, xcenter = 5000,ycenter = 5000, img = None, psf = None):
    sed = generateSEDfromFiles(sedFile,2000,10000,0.5) # in photo/s/m2/A

    throughput = Table.read(throughputFn)
    gal = galsim.Sersic(s_n, half_light_radius=re)

    gal_pa = pa * galsim.degrees
    gal_ell = gal.shear(q=ell, beta=gal_pa)

    conv_gal = galsim.Convolve([gal_ell,psf])

    # the throughput curve contains the information to convert photos to electrons
    flux = intCSSTFlux(spec = sed, throughput=throughput) # unit: e-/s

    #       ==== normalized image profile (sum to 1) =======
    stamp = conv_gal.drawImage(wcs=galsim.PixelScale(p_size))*t*flux

    stamp.setCenter(xcenter,ycenter)

    bounds = stamp.bounds & img.bounds

    img[bounds] = img[bounds] + stamp[bounds]

p_size = 0.074
# psf = galsim.Moffat(beta=100, fwhm=0.39)
psf = galsim.Gaussian(fwhm=0.39)

s_n = 1.0
re = 0.1
pa = 90
ell = 0.6
t = 150
aper = 2
skybg = 0.3; ## e-/pixel/s
xcenter = 1500
ycenter = 500
nx = 3000
ny = 1000

sedFile = '../data/sed/sed_90156.txt'
#import ObservationSim
#DATAROOT = Path(ObservationSim.__file__).parent / "Instrument" / "data" / "sls_conf"
#specConfile = DATAROOT / 'CSST_GI2.conf'
specConfile = '../data/conf/CSST_GI2.conf'

# img = galsim.ImageF(9232, 9216)
img = galsim.ImageF(nx, ny)

generateSpec(sedFile = sedFile, specConfile = specConfile,s_n = s_n, re = re, pa = pa, ell = ell, t = t, aper = aper, xcenter = xcenter,ycenter = ycenter, img = img, psf = psf)

img_a = img.array + skybg*t

img_a = np.random.poisson(img_a)
img_a = addReadoutNois(img = img_a, readout = 5)

fits.writeto('test_spec.fits',img_a,overwrite=True)

img = galsim.ImageF(nx, ny)
througputFn = '../data/CSST_Throughput/i.Throughput.fits'
generateGal(sedFile = sedFile, throughputFn = througputFn,s_n = s_n, re = re, pa = pa, ell = ell, t = t, aper = aper, xcenter = xcenter,ycenter = ycenter, img = img, psf = psf)


img_a = img.array + skybg*t

img_a = np.random.poisson(img_a)
img_a = addReadoutNois(img = img_a, readout = 5)

fits.writeto('test_img.fits',img_a,overwrite=True)
